package cn.kinoko.common.utils;

import lombok.Getter;
import lombok.Setter;

import java.util.Arrays;

/**
 * 线性回归工具
 * @author Kinoko
 * @date 2023/10/26
 */
public class LinearRegression {

    private final double[] xData;
    private final double[] yData;
    @Getter
    @Setter
    private double slope;
    @Getter
    @Setter
    private double intercept;

    public LinearRegression(double[] xData, double[] yData) {
        this.xData = xData;
        this.yData = yData;
    }

    public void fit() {
        int n = xData.length;
        double sumX = Arrays.stream(xData).sum();
        double sumY = Arrays.stream(yData).sum();
        double sumXY = 0;
        double sumXSquare = 0;

        for (int i = 0; i < n; i++) {
            sumXY += xData[i] * yData[i];
            sumXSquare += xData[i] * xData[i];
        }

        slope = (n * sumXY - sumX * sumY) / (n * sumXSquare - sumX * sumX);
        intercept = (sumY - slope * sumX) / n;
    }

    /**
     * 更新斜率和截距。
     *
     * @param oldSlope 当前斜率
     * @param oldIntercept 当前截距
     * @param x 新的数据点特征值数组
     * @param y 新的数据点标签值数组
     * @param learningRate 学习率
     * @return 更新后的斜率和截距
     */
    public static double[] updateParameters(double oldSlope, double oldIntercept, double[] x, double[] y, double learningRate) {
        double newSlope = oldSlope;
        double newIntercept = oldIntercept;

        // 验证 x 和 y 数组长度一致
        if (x.length != y.length) {
            throw new IllegalArgumentException("x and y must have the same length");
        }

        // 计算梯度并更新参数
        for (int i = 0; i < x.length; i++) {
            // 计算预测值
            double prediction = oldSlope * x[i] + oldIntercept;
            // 计算误差
            double error = y[i] - prediction;

            // 更新斜率
            newSlope += (learningRate / x.length) * error * x[i];
            // 更新截距
            newIntercept += (learningRate / x.length) * error;
        }

        return new double[]{newSlope, newIntercept};
    }



    public double predict(double x) {
        return slope * x + intercept;
    }

    public static double predict(double x, double slope, double intercept) {
        return slope * x + intercept;
    }

    public static void main(String[] args) {
        // 原始坐标和目标坐标对
        double[] xValues = {483, 233, 1472, 1367, 710, 441, 222, 1574, 1359, 1159, 1164, 926, 552, 238, 1544, 1117, 1084, 1122, 809, 683, 402}; // 原始 x 坐标
        double[] xPrimeValues = {293, 248, 491, 467, 341, 286, 246, 505, 467, 429, 428, 380, 308, 249, 499, 419, 416, 418, 361, 338, 281}; // 目标 x 坐标

        double[] yValues = {233, 838, 1407, 2231, 1418, 1414, 1939, 229, 304, 331, 506, 249, 253, 676, 1286, 1203, 1770, 2196, 1226, 1757, 1250, 1531}; // 原始 y 坐标
        double[] yPrimeValues = {36, 151, 263, 417, 263, 265, 363, 34, 47, 55, 89, 37, 42, 118, 238, 222, 328, 408, 229, 331, 232, 284}; // 目标 y 坐标

        LinearRegression regressionX = new LinearRegression(xValues, xPrimeValues);
        LinearRegression regressionY = new LinearRegression(yValues, yPrimeValues);
        regressionX.fit();
        regressionY.fit();

        System.out.println("XSlope: " + regressionX.slope);
        System.out.println("XIntercept: " + regressionX.intercept);
        System.out.println("YSlope: " + regressionY.slope);
        System.out.println("YIntercept: " + regressionY.intercept);

        double xToPredict = 800;
        double xPrediction = regressionX.predict(xToPredict);
        System.out.println("Prediction for x = " + xToPredict + ": " + xPrediction);

        double yToPredict = 2500;
        double yPrediction = regressionY.predict(yToPredict);
        System.out.println("Prediction for y = " + yToPredict + ": " + yPrediction);

        System.out.println("增量测试...");
        double learningRate = 0.00001;
        double[] oldX = {196};
        double[] newX = {244};

        double[] parameters = updateParameters(regressionX.slope, regressionX.intercept, oldX, newX, learningRate);
        regressionX.setSlope(parameters[0]);
        regressionX.setIntercept(parameters[1]);
        System.out.println("Updated XSlope: " + regressionX.getSlope());
        System.out.println("Updated XIntercept: " + regressionX.getIntercept());

        double updateXPrediction = regressionX.predict(xToPredict);
        System.out.println("Updated Prediction for x = " + xToPredict + ": " + updateXPrediction);
        double updateYPrediction = regressionY.predict(yToPredict);
        System.out.println("Updated Prediction for y = " + yToPredict + ": " + updateYPrediction);
    }
}